Ruihang Lai
8813d0a2bd
## Context
When dealing with end-to-end models, we note that some tensors may have large shapes. Thus, when designing graph-level IR, we sometimes use `int64` instead of `int32` for the shape. Below is an dense GeMM example which has `int64` input tensor shape:
```python
@tvm.script.ir_module
class Module:
@T.prim_func
def main(rxplaceholder: T.Buffer[(1, 512), "float32"], rxplaceholder_1: T.Buffer[(T.int64(1000), T.int64(512)), "float32"], T_matmul_NT: T.Buffer[(1, T.int64(1000)), "float32"]) -> None:
# function attr dict
T.func_attr({"global_symbol": "dense", "tir.noalias": True, "op_pattern": 3})
# body
# with T.block("root")
for i0_0, i1_0, i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(1, 4, 1, 25, 8, 1, 10, 64, 1, 1):
with T.block("T_matmul_NT"):
i = T.axis.spatial(1, 0)
j = T.axis.spatial(T.int64(1000), i1_0 * T.int64(250) + i1_1 * T.int64(10) + i1_2)
k = T.axis.reduce(512, i2_0 * 64 + i2_1)
T.reads(T_matmul_NT[i, j], rxplaceholder[i, k], rxplaceholder_1[j, k])
T.writes(T_matmul_NT[i, j])
T.block_attr({"layout_free_placeholders":[rxplaceholder_1], "meta_schedule.tiling_structure":"SSRSRS"})
with T.init():
T_matmul_NT[i, j] = T.float32(0)
T_matmul_NT[i, j] = T_matmul_NT[i, j] + rxplaceholder[i, k] * rxplaceholder_1[j, k]
```
## Problem
Though our TVMScript printer can easily print `int64` constants, the parser had poor support for `int64`. So this PR introduces some parser support for `int64`, basically about the data type of loop variables, block iterators and block read/write regions.
Besides the parser, most of the TIR schedule primitives didn't take `int64` into account in their implementations. These schedule primitives will be fixed and updated in recent future, in followup PRs.
56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you 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.
|
|
"""Helper functions in TVM Script Parser"""
|
|
|
|
from typing import List, Optional
|
|
|
|
from tvm.arith import Analyzer
|
|
from tvm.ir import Range
|
|
from tvm.tir import PrimExpr, BufferRegion
|
|
from tvm.tir.expr import IntImm
|
|
from .node import BufferSlice
|
|
|
|
|
|
def buffer_slice_to_region(
|
|
buffer_slice: BufferSlice, analyzer: Optional[Analyzer] = None
|
|
) -> BufferRegion:
|
|
"""Construct BufferRegion from BufferSlice
|
|
|
|
Parameters
|
|
----------
|
|
buffer_slice : BufferSlice
|
|
The input BufferSlice
|
|
|
|
analyzer : Optional[tvm.arith.Analyzer]
|
|
The analyzer for simplifying. If not provided, the method will construct a new one
|
|
|
|
Returns
|
|
-------
|
|
buffer_region : BufferRegion
|
|
The constructed BufferRegion.
|
|
"""
|
|
region: List[Range] = []
|
|
for s in buffer_slice.slices:
|
|
start = s.start if isinstance(s.start, PrimExpr) else IntImm("int32", s.start)
|
|
extent = IntImm(start.dtype, 1) if s.stop is None else s.stop - s.start
|
|
if not analyzer:
|
|
analyzer = Analyzer()
|
|
if isinstance(extent, PrimExpr):
|
|
extent = analyzer.simplify(extent)
|
|
region.append(Range.from_min_extent(start, extent, span=s.span))
|
|
return BufferRegion(buffer_slice.buffer, region)
|