[Hidet script] Add hidet.lang.types submodule (#340)

Also add copyright information for some files.

python/hidet/apps/__init__.py

@@ -0,0 +1,11 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

python/hidet/apps/compile_server/__init__.py

@@ -1,2 +1,13 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from .compilation import remote_build
 from .core import init_api

python/hidet/apps/compile_server/auth.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import requests
 from .core import api_url
 

python/hidet/apps/compile_server/compilation.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import zipfile
 import shutil
 import tempfile

python/hidet/apps/compile_server/core.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import Optional
 import hidet
 

python/hidet/apps/compile_server/user.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import requests
 from .core import access_token, api_url
 

python/hidet/cuda/capability.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import Optional
 from dataclasses import dataclass
 

python/hidet/graph/ops/conv1d/resolve.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import List, Optional
 from hidet.graph.operator import Operator, Tensor
 from hidet.graph.transforms import ResolveRule, register_resolve_rule

python/hidet/graph/ops/matmul/batch_matmul.py

@@ -14,7 +14,7 @@
 from hidet.ir import IRModule
 from hidet.ir.compute import reduce
 from hidet.ir.expr import is_constant, cast
-from hidet.ir.layout import StridesLayout, data_layout, row_major, column_major, local_layout
+from hidet.ir.layout import StridesLayout, strided_layout, row_major, column_major, local_layout
 from hidet.ir.type import data_type, TensorType, DataType, void_p
 from hidet.lang import i32, spatial, repeat, register_tensor, shared_tensor, attrs, grid, tensor_pointer
 from hidet.lang.cuda import blockIdx, threadIdx, syncthreads
@@ -437,9 +437,9 @@ def resolve_mma_type(a_dtype: DataType, b_dtype: DataType, c_dtype: DataType):
             task_shape=[block_k, block_n], num_workers=num_threads, ranks=[0, 1]
         )
 
-        smem_a_layout = data_layout([2, block_m, block_k], ranks=[0, 1, 2])
-        smem_b_layout = data_layout([2, block_k, block_n], ranks=[0, 1, 2])
-        smem_c_layout = data_layout([block_m, block_n], ranks=[0, 1])
+        smem_a_layout = strided_layout([2, block_m, block_k], ranks=[0, 1, 2])
+        smem_b_layout = strided_layout([2, block_k, block_n], ranks=[0, 1, 2])
+        smem_c_layout = strided_layout([block_m, block_n], ranks=[0, 1])
         regs_a_layout = row_major(2, mma_count_m, mma_config.a_elements)
         regs_b_layout = row_major(2, mma_count_n, mma_config.b_elements)
         regs_c_layout = row_major(mma_count_m, mma_count_n, mma_config.c_elements)

python/hidet/graph/transforms/graph_patterns/quant/__init__.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import List, Union
 from ..base import SubgraphRewriteRule
 

python/hidet/graph/transforms/graph_patterns/register_all_patterns.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 # pylint: disable=unused-import
 from .arithmetic_patterns import arithmetic_patterns
 from .transform_patterns import transform_patterns

python/hidet/graph/transforms/selective_quantize.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import List
 
 from hidet.graph.flow_graph import FlowGraph

python/hidet/ir/expr.py

@@ -10,7 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # pylint: disable=import-outside-toplevel, useless-parent-delegation, redefined-outer-name, redefined-builtin
-# pylint: disable=useless-super-delegation
+# pylint: disable=useless-super-delegation, protected-access
 from __future__ import annotations
 from typing import Optional, Union, Sequence, Tuple, Dict, Type, Callable
 import string
@@ -815,6 +815,11 @@ def logical_not(a: Union[Expr, PyScalar]):
     return Expr._unary(LogicalNot, a)
 
 
+def bitwise_not(a: Union[Expr, PyScalar]):
+    a = convert(a)
+    return Expr._unary(BitwiseNot, a)
+
+
 def equal(a: Union[Expr, PyScalar], b: Union[Expr, PyScalar]):
     a = convert(a)
     b = convert(b)

python/hidet/ir/layout.py

@@ -325,7 +325,7 @@ def local_layout(*shape: Int):
     return LocalLayout(shape)
 
 
-def data_layout(shape: Sequence[Int], ranks: Optional[List[int]] = None):
+def strided_layout(shape: Sequence[Int], ranks: Optional[List[int]] = None):
     if ranks is None:
         ranks = list(range(len(shape)))
     return StridesLayout.from_shape(shape, ranks)

python/hidet/ir/library/__init__.py

@@ -1 +1,12 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from . import cuda

python/hidet/ir/library/cuda/__init__.py

@@ -1 +1,12 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from .matmul import matmul_simt

python/hidet/ir/library/cuda/matmul/__init__.py

@@ -1 +1,12 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from .simt import matmul_simt

python/hidet/ir/library/cuda/matmul/simt.py

@@ -1,10 +1,21 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import List
 
 import hidet
 from hidet.ir.type import DataType, TensorType, tensor_type, tensor_pointer_type
 from hidet.ir.dtypes import vectorize, i32
 from hidet.ir.expr import Expr, is_true
-from hidet.ir.layout import row_major, data_layout, local_layout
+from hidet.ir.layout import row_major, strided_layout, local_layout
 from hidet.ir.library import tune
 from hidet.ir.library.utils import get_tensor_type
 from hidet.ir.primitives.runtime import request_cuda_workspace
@@ -157,8 +168,8 @@ def matmul_simt(
     # prepare data layout
     tune.check(block_k % lanes == 0)
     block_k_vectors = block_k // lanes
-    smem_a_layout = data_layout([2, block_k_vectors, block_m + 1])
-    smem_b_layout = data_layout([2, block_k_vectors, block_n + 1])
+    smem_a_layout = strided_layout([2, block_k_vectors, block_m + 1])
+    smem_b_layout = strided_layout([2, block_k_vectors, block_n + 1])
     regs_a_layout = (  # 2 x block_m
         row_major(2, 1)
         .local(1, block_warps[0])

python/hidet/ir/library/utils.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from hidet.ir.tools import infer_type
 from hidet.ir.type import TensorType, TensorPointerType
 from hidet.ir.expr import Expr

python/hidet/ir/primitives/cuda/errchk.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from hidet.ir.stmt import BlackBoxStmt
 
 

python/hidet/ir/target.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from __future__ import annotations
 from typing import List, Dict
 

python/hidet/ir/utils/broadcast_utils.py

@@ -1,3 +1,14 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import Sequence, List
 from hidet.ir.expr import Expr, Int, is_constant, if_then_else
 

python/hidet/lang/__init__.py

@@ -14,7 +14,7 @@
 from hidet.lang.constructs import meta
 
 from hidet.ir.type import BaseType, DataType, TensorType, PointerType, VoidType, ReferenceType, void_p, data_type
-from hidet.ir.expr import Expr, Var, cast, view, address, Dereference
+from hidet.ir.expr import Expr, Var, cast, view, address, Dereference, bitwise_not
 from hidet.ir.mapping import row_spatial, row_repeat, col_repeat, col_spatial, TaskMapping, auto_map
 from hidet.ir.layout import DataLayout
 from hidet.ir.primitives import printf

python/hidet/lang/constructs/loops.py

@@ -128,7 +128,7 @@ def grid(*dim_extents, attrs: Optional[str] = None, bind_tuple=False):
 
     Parameters
     ----------
-    dim_extents: Sequence[Expr or int or list or tuple or str]
+    dim_extents: Expr or int or list or tuple or str
         The length of each dimension. The last one can be the attrs.
 
     attrs: Optional[str]

python/hidet/lang/layout.py

@@ -10,4 +10,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # pylint: disable=unused-import
-from hidet.ir.layout import row_major, column_major, local_layout, DataLayout
+from hidet.ir.layout import row_major, column_major, local_layout, strided_layout, DataLayout

python/hidet/lang/types.py

@@ -0,0 +1,19 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# pylint: disable=unused-import
+from hidet.ir.dtypes import i8, i16, i32, i64, u8, u16, u32, u64, f16, f32, f64, bf16, tf32
+from hidet.ir.dtypes import int8, int16, int32, int64, uint8, uint32, uint64, float16, float32, float64, bfloat16
+from hidet.ir.dtypes import tfloat32
+
+from hidet.ir.type import void_p, void, byte_p
+
+from hidet.lang.constructs.declare import register_tensor, shared_tensor, tensor_pointer, tensor, DeclareScope