    def _compute_workspace_size(
        self,
        shape: List[IntVar],
        reduction_axis: int,
        input_type: str,
    ) -> int:
        """
        Compute workspace size for the given shape using the same algorithm as
        cutlass's TensorReduction kernel. The only difference is that we use
        the maximum dim value for dynamic dimension, whereas TensorReduction
        uses the real dim value at runtime.
        """
        # workspace size must be computed in terms of the accumulation dtype;
        # see the CUTLASS code in TensorReductionAffineStrided for the reference:
        # https://github.com/NVIDIA/cutlass/blob/ff02da266713bd3365aed65c552412e126c040cb/include/cutlass/reduction/device/tensor_reduce_affine_strided.h#L223
        accumulation_type = "float32"
        try:
            if (
                backend.target.Target.current()._kwargs.get("use_fp16_acc", False)
                and input_type == "float16"
            ):
                accumulation_type = input_type
        except RuntimeError:
            # Target is not set: conservatively
            # assume float32 accumulation type
            pass
        # Make sure the last dim is static to pre-compute vector_length.
        # Note that this is a temporary constraint. Once we replace TensorReduction
        # with our own col-reduction kernel, we will remove this entire workaround.
        if not shape_utils.is_static_dimension(shape, -1):
            raise NotImplementedError("expect the last dim to be static")
        last_dim = shape[-1].value()
        vector_lens_config = [32, 16, 8, 4, 1]
        vector_length = 1
        for vec_len in vector_lens_config:
            if last_dim % vec_len == 0:
                vector_length = vec_len
                break
        num_dims = 4
        rank = len(shape)
        assert (
            rank <= num_dims
        ), f"expected rank <= num_dims, but got {rank=}, {num_dims=}"
        # adjust reduction axis
        reduction_axis = num_dims - rank + reduction_axis
        all_shape_dims = [
            list(range(d.lower_bound(), d.upper_bound() + 1)) for d in shape
        ]
        max_ws = 0
        # Go through cartesian product of all possible dynamic dim values
        # to find the maximum workspace size. It might be a bit heavy
        # for some cases. However, it's OK for our current use cases where
        # we have a single dynamic axis within a range of several thousand.
        # Moreover, we would remove this entire estimation once we have
        # our own row-reduction kernel.
        for one_dims in itertools.product(*all_shape_dims):
            extent_affine = []
            prefix_dims = [1] * (num_dims - rank)
            # Use dim's upper-bound for computing workspace size
            shape_dims = prefix_dims + list(one_dims)
            # normalize extent_affine list
            # reference: TensorReduction construction in cutlass
            # reduction/device/tensor_reduce.h
            if reduction_axis == 0:
                extent_affine.append(shape_dims[1])
                extent_affine.append(shape_dims[2])
                extent_affine.append(shape_dims[0])
                extent_affine.append(shape_dims[3])
            elif reduction_axis == 1:
                extent_affine.append(shape_dims[0])
                extent_affine.append(shape_dims[2])
                extent_affine.append(shape_dims[1])
                extent_affine.append(shape_dims[3])
            elif reduction_axis == 2:
                extent_affine.append(shape_dims[0])
                extent_affine.append(shape_dims[1])
                extent_affine.append(shape_dims[2])
                extent_affine.append(shape_dims[3])
            else:
                # note that we already ruled out non-col-reduction kernels so that
                # reduction_axis would never be 3. Consequently, we would never
                # invoke contiguous tensor_reduce kernels.
                raise RuntimeError(
                    f"Expected reduction_axis to be within [0, 2], but got {reduction_axis=}"
                )
            max_ws = max(
                max_ws,
                self._compute_ws_size_strided(
                    extent_affine,
                    reduction_axis,
                    vector_length,
                    accumulation_type,
                ),
            )
        return max_ws
Compute workspace size for the given shape using the same algorithm as
cutlass's TensorReduction kernel. The only difference is that we use
the maximum dim value for dynamic dimension, whereas TensorReduction
uses the real dim value at runtime.
python/aitemplate/compiler/ops/reduce/reduce_common.py